package com.jd.genie.agent.llm;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.TypeReference;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jd.genie.agent.agent.AgentContext;
import com.jd.genie.agent.dto.Message;
import com.jd.genie.agent.dto.tool.McpToolInfo;
import com.jd.genie.agent.dto.tool.ToolCall;
import com.jd.genie.agent.dto.tool.ToolChoice;
import com.jd.genie.agent.printer.Printer;
import com.jd.genie.agent.tool.BaseTool;
import com.jd.genie.agent.tool.ToolCollection;
import com.jd.genie.agent.util.SpringContextHolder;
import com.jd.genie.agent.util.StringUtil;
import com.jd.genie.config.GenieConfig;
import com.jd.genie.config.LLMConfig;
import lombok.*;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.apache.commons.lang3.StringUtils;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * LLM 类
 */
@Slf4j
@Data
public class LLM {

    private static final Map<String, LLM> instances = new ConcurrentHashMap<>();

    private final String model;
    private final String llmErp;
    private final int maxTokens;
    private final double temperature;
    private final String apiKey;
    private final String baseUrl;
    private final String interfaceUrl;
    private final TokenCounter tokenCounter;
    private final ObjectMapper objectMapper;
    private final Map<String, Object> extParams;

    private int totalInputTokens;
    private Integer maxInputTokens;

    public LLM(String modelName, String llmErp)
    {
        this.llmErp = llmErp;

        LLMSettings config = LLMConfig.getLLMConfig(modelName);
        this.model = config.getModel();
        this.maxTokens = config.getMaxTokens();
        this.temperature = config.getTemperature();
        this.apiKey = config.getApiKey();
        this.baseUrl = config.getBaseUrl();
        this.interfaceUrl = config.getInterfaceUrl();

        // 初始化 token 计数相关属性
        this.totalInputTokens = 0;
        this.maxInputTokens = config.getMaxInputTokens();
        this.extParams = config.getExtParams();

        // 初始化 tokenizer
        this.tokenCounter = new TokenCounter();
        this.objectMapper = new ObjectMapper();
    }

    /**
     * 格式化消息
     */
    public static List<Map<String, Object>> formatMessages(List<Message> messages)
    {
        List<Map<String, Object>> formattedMessages = new ArrayList<>();

        for (Message message : messages) {
            Map<String, Object> messageMap = new HashMap<>();
            // 处理 base64 图像
            if (message.getBase64Image() != null && !message.getBase64Image().isEmpty()) {
                List<Map<String, Object>> multimodalContent = new ArrayList<>();
                // 创建内层的 image_url Map
                Map<String, String> imageUrlMap = new HashMap<>();
                imageUrlMap.put("url", "data:image/jpeg;base64," + message.getBase64Image());
                // 创建外层的 Map
                Map<String, Object> outerMap = new HashMap<>();
                outerMap.put("type", "image_url");
                outerMap.put("image_url", imageUrlMap);
                // 将创建好的 Map 添加到 multimodalContent 中
                multimodalContent.add(outerMap);

                Map<String, Object> contentMap = new HashMap<>();
                outerMap.put("type", "text");
                outerMap.put("text", message.getContent());
                multimodalContent.add(contentMap);

                messageMap.put("role", message.getRole().getValue());
                messageMap.put("content", multimodalContent);
            } else if (message.getToolCalls() != null && !message.getToolCalls().isEmpty()) {
                messageMap.put("role", message.getRole().getValue());
                List<Map<String, Object>> toolCallsMap = JSON.parseObject(JSON.toJSONString(message.getToolCalls()), new TypeReference<>() {
                });
                messageMap.put("tool_calls", toolCallsMap);
            } else if (message.getToolCallId() != null && !message.getToolCallId().isEmpty()) {
                GenieConfig genieConfig = SpringContextHolder.getApplicationContext().getBean(GenieConfig.class);
                // 脱敏
                String content = StringUtil.textDesensitization(message.getContent(), genieConfig.getSensitivePatterns());
                messageMap.put("role", message.getRole().getValue());
                messageMap.put("content", content);
                messageMap.put("tool_call_id", message.getToolCallId());
            } else {
                messageMap.put("role", message.getRole().getValue());
                messageMap.put("content", message.getContent());
            }

            formattedMessages.add(messageMap);
        }

        return formattedMessages;
    }

    /**
     * 向 LLM 发送请求并获取响应
     */
    public String ask(
            AgentContext context,
            List<Message> messages,
            List<Message> systemMsgs,
            boolean stream,
            Double temperature
    )
    {
        List<Map<String, Object>> formattedMessages;
        // 格式化系统和用户消息
        if (systemMsgs != null && !systemMsgs.isEmpty()) {
            List<Map<String, Object>> formattedSystemMsgs = formatMessages(systemMsgs);
            formattedMessages = new ArrayList<>(formattedSystemMsgs);
            formattedMessages.addAll(formatMessages(messages));
        } else {
            formattedMessages = formatMessages(messages);
        }

        // 准备请求参数
        Map<String, Object> params = new HashMap<>();
        params.put("model", model);
        if (StringUtils.isNotEmpty(llmErp)) {
            params.put("erp", llmErp);
        }
        params.put("messages", formattedMessages);

        // 根据模型设置不同的参数
        params.put("max_tokens", maxTokens);
        params.put("temperature", temperature != null ? temperature : this.temperature);
        if (Objects.nonNull(extParams)) {
            params.putAll(extParams);
        }

        log.info("{} call llm ask request {}", context.getRequestId(), JSONObject.toJSONString(params));
        // 处理非流式请求
        if (!stream) {
            params.put("stream", false);

            // 调用 API
            String response = callOpenAI(params);
            try {
                // 解析响应
                log.info("{} call llm response {}", context.getRequestId(), response);
                JsonNode jsonResponse = objectMapper.readTree(response);
                JsonNode choices = jsonResponse.get("choices");

                if (choices == null || choices.isEmpty() || choices.get(0).get("message").get("content") == null) {
                    throw new IllegalArgumentException("Empty or invalid response from LLM");
                }

                return choices.get(0).get("message").get("content").asText();
            } catch (IOException e) {
                throw new CompletionException(e);
            }
        } else {
            // 处理流式请求
            params.put("stream", true);
            // 调用流式 API
            return callOpenAIStream(params);
        }
    }

    /**
     * 向 LLM 发送工具请求并获取响应
     */
    public ToolCallResponse askTool(
            AgentContext context,
            List<Message> messages,
            Message systemMsgs,
            ToolCollection tools,
            ToolChoice toolChoice,
            Double temperature,
            boolean stream,
            int timeout,
            Printer printer)
    {
        // 验证 toolChoice
        if (!ToolChoice.isValid(toolChoice)) {
            throw new IllegalArgumentException("Invalid tool_choice: " + toolChoice);
        }
        long startTime = System.currentTimeMillis();

        // 设置 API 请求
        Map<String, Object> params = new HashMap<>();

        // tools
        List<Map<String, Object>> formattedTools = new ArrayList<>();
        // base tool
        for (BaseTool tool : tools.getToolMap().values()) {
            Map<String, Object> functionMap = new HashMap<>();
            functionMap.put("name", tool.getName());
            functionMap.put("description", tool.getDescription());
            functionMap.put("parameters", tool.toParams());
            Map<String, Object> toolMap = new HashMap<>();
            toolMap.put("type", "function");
            toolMap.put("function", functionMap);
            formattedTools.add(toolMap);
        }
        // mcp tool
        for (McpToolInfo tool : tools.getMcpToolMap().values()) {
            Map<String, Object> parameters = JSON.parseObject(tool.getParameters(), new TypeReference<Map<String, Object>>() {
            });
            Map<String, Object> functionMap = new HashMap<>();
            functionMap.put("name", tool.getName());
            functionMap.put("description", tool.getDesc());
            functionMap.put("parameters", parameters);
            Map<String, Object> toolMap = new HashMap<>();
            toolMap.put("type", "function");
            toolMap.put("function", functionMap);
            formattedTools.add(toolMap);
        }

        // 格式化消息
        List<Map<String, Object>> formattedMessages = new ArrayList<>();
        if (Objects.nonNull(systemMsgs)) {
            formattedMessages.addAll(formatMessages(List.of(systemMsgs)));
        }
        formattedMessages.addAll(formatMessages(messages));

        params.put("model", model);
        if (StringUtils.isNotEmpty(llmErp)) {
            params.put("erp", llmErp);
        }
        params.put("messages", formattedMessages);

        params.put("tools", formattedTools);
        params.put("tool_choice", toolChoice.getValue());

        // 添加模型特定参数
        params.put("max_tokens", maxTokens);
        params.put("temperature", temperature != null ? temperature : this.temperature);
        if (Objects.nonNull(extParams)) {
            params.putAll(extParams);
        }

        log.info("{} call llm request {}", context.getRequestId(), JSONObject.toJSONString(params));
        if (!stream) {
            params.put("stream", false);
            // 调用 API
            String responseJson = callOpenAI(params, timeout);
            try {
                // 解析响应
                log.info("{} call llm response {}", context.getRequestId(), responseJson);
                JsonNode jsonResponse = objectMapper.readTree(responseJson);
                JsonNode choices = jsonResponse.get("choices");

                if (choices == null || choices.isEmpty() || choices.get(0).get("message") == null) {
                    log.error("{} Invalid response: {}", context.getRequestId(), responseJson);
                    throw new IllegalArgumentException("Invalid or empty response from LLM");
                }

                // 提取响应内容
                JsonNode message = choices.get(0).get("message");
                String content = message.has("content") && !"null".equals(message.get("content").asText()) ? message.get("content").asText() : null;

                // 提取工具调用
                List<ToolCall> toolCalls = new ArrayList<>();
                if (message.has("tool_calls")) {
                    JsonNode toolCallsNode = message.get("tool_calls");
                    for (JsonNode toolCall : toolCallsNode) {
                        String id = toolCall.get("id").asText();
                        String type = toolCall.get("type").asText();

                        // 提取函数信息
                        JsonNode functionNode = toolCall.get("function");
                        String name = functionNode.get("name").asText();
                        String arguments = functionNode.get("arguments").asText();
                        toolCalls.add(new ToolCall(id, type, new ToolCall.Function(name, arguments)));
                    }
                }
                // 提取其他信息
                String finishReason = choices.get(0).get("finish_reason").asText();
                int totalTokens = jsonResponse.get("usage").get("total_tokens").asInt();

                long endTime = System.currentTimeMillis();
                long duration = endTime - startTime;
                return new ToolCallResponse(content, toolCalls, finishReason, totalTokens, duration);
            } catch (IOException e) {
                throw new CompletionException(e);
            }
        } else {
            // 处理流式请求
            params.put("stream", true);
            // 调用流式 API
            return callOpenAIFunctionCallStream(context, params, printer);
        }

    }

    /**
     * 调用 OpenAI API（抽象方法，实际实现需要在子类中提供）
     */
    protected String callOpenAI(Map<String, Object> params)
    {
        return callOpenAI(params, 300); // 默认超时时间为 300 秒
    }

    /**
     * 调用 OpenAI API（抽象方法，实际实现需要在子类中提供）
     */
    @SneakyThrows
    protected String callOpenAI(Map<String, Object> params, int timeout)
    {
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(timeout, TimeUnit.SECONDS)
                .readTimeout(timeout, TimeUnit.SECONDS)
                .writeTimeout(timeout, TimeUnit.SECONDS)
                .build();

        String apiEndpoint = baseUrl + interfaceUrl;

        RequestBody body = RequestBody.create(
                MediaType.parse("application/json"),
                objectMapper.writeValueAsString(params)
        );

        Request.Builder requestBuilder = new Request.Builder()
                .url(apiEndpoint)
                .post(body);

        // 添加适当的认证头
        requestBuilder.addHeader("Authorization", "Bearer " + apiKey);

        Request request = requestBuilder.build();
        try (Response response = client.newCall(request).execute();
             ResponseBody responseBody = response.body()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected response code: " + response);
            } else {
                return responseBody.string();
            }
        }
    }

    /**
     * 调用 OpenAI 流式 API（抽象方法，实际实现需要在子类中提供）
     */
    @SneakyThrows
    public ToolCallResponse callOpenAIFunctionCallStream(AgentContext context, Map<String, Object> params, Printer printer)
    {
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(300, TimeUnit.SECONDS)
                .readTimeout(300, TimeUnit.SECONDS)
                .writeTimeout(300, TimeUnit.SECONDS)
                .build();

        String apiEndpoint = baseUrl + interfaceUrl;
        RequestBody body = RequestBody.create(
                MediaType.parse("application/json"),
                objectMapper.writeValueAsString(params)
        );
        Request.Builder requestBuilder = new Request.Builder()
                .url(apiEndpoint)
                .post(body);
        // 添加适当的认证头
        requestBuilder.addHeader("Authorization", "Bearer " + apiKey);
        Request request = requestBuilder.build();

        final int firstInterval = 1;
        final int sendInterval = 3;

        try (Response response = client.newCall(request).execute();
             ResponseBody responseBody = response.body()) {
            boolean isContent = true;

            if (!response.isSuccessful() || responseBody == null) {
                log.error("{} ask tool stream response error or empty", context.getRequestId());
                throw new IOException("Unexpected response code: " + response);
            }

            String messageId = StringUtil.getUUID();
            StringBuilder stringBuilder = new StringBuilder();
            StringBuilder stringBuilderAll = new StringBuilder();
            int index = 1;
            Map<Integer, OpenAIToolCall> openToolCallsMap = new HashMap<>();
            String line;
            BufferedReader reader = new BufferedReader(new InputStreamReader(responseBody.byteStream()));
            while ((line = reader.readLine()) != null) {
                if (line.startsWith("data: ")) {
                    String data = line.substring(6);
                    if (data.equals("[DONE]")) {
                        break;
                    }
                    try {
                        JsonNode chunk = objectMapper.readTree(data);
                        if (chunk.has("choices") && !chunk.get("choices").isEmpty()) {
                            for (JsonNode element : chunk.get("choices")) {
                                OpenAIChoice choice = objectMapper.convertValue(element, OpenAIChoice.class);
                                // content
                                if (Objects.nonNull(choice.delta.content)) {
                                    String content = choice.delta.content;
                                    // log.info("{} recv content data: >>{}<<", context.getRequestId(), content);
                                    if (!isContent) { // 忽略json内容
                                        stringBuilderAll.append(content);
                                        continue;
                                    }
                                    stringBuilder.append(content);
                                    stringBuilderAll.append(content);
                                    if (index == firstInterval || index % sendInterval == 0) {
                                        printer.send(messageId, context.getStreamMessageType(), stringBuilder.toString(), null, false);
                                        stringBuilder.setLength(0);
                                    }
                                    index++;
                                }
                                // tool call
                                if (Objects.nonNull(choice.delta.tool_calls)) {
                                    List<OpenAIToolCall> openAIToolCalls = choice.delta.tool_calls;
                                    // log.info("{} recv tool call data: {}", context.getRequestId(), openAIToolCalls);
                                    for (OpenAIToolCall toolCall : openAIToolCalls) {
                                        OpenAIToolCall currentToolCall = openToolCallsMap.get(toolCall.index);
                                        if (Objects.isNull(currentToolCall)) {
                                            currentToolCall = new OpenAIToolCall();
                                        }
                                        // [{"index":0,"id":"call_j74R8JMFWTC4rW5wHJ0TtmNU","type":"function","function":{"name":"planning","arguments":""}}]
                                        if (Objects.nonNull(toolCall.id)) {
                                            currentToolCall.id = toolCall.id;
                                        }
                                        if (Objects.nonNull(toolCall.type)) {
                                            currentToolCall.type = toolCall.type;
                                        }
                                        if (Objects.nonNull(toolCall.function)) {
                                            if (Objects.nonNull(toolCall.function.name)) {
                                                currentToolCall.function = toolCall.function;
                                            }
                                            if (Objects.nonNull(toolCall.function.arguments)) {
                                                if (currentToolCall.function.arguments == null) {
                                                    currentToolCall.function.arguments = toolCall.function.arguments;
                                                } else {
                                                    currentToolCall.function.arguments += toolCall.function.arguments;
                                                }
                                            }
                                        }
                                        openToolCallsMap.put(toolCall.index, currentToolCall);
                                    }
                                }
                            }
                        }
                    } catch (Exception e) {
                        log.error("{} process response error", context.getRequestId(), e);
                    }
                }
            }

            String contentAll = stringBuilderAll.toString();
            // function_call
            if (!contentAll.isEmpty()) {
                printer.send(messageId, context.getStreamMessageType(), stringBuilder.toString(), null, false);
                printer.send(messageId, context.getStreamMessageType(), stringBuilderAll.toString(), null, true);
            }

            List<ToolCall> toolCalls = new ArrayList<>();
            // function call
            for (OpenAIToolCall toolCall : openToolCallsMap.values()) {
                toolCalls.add(ToolCall.builder()
                        .id(toolCall.id)
                        .type(toolCall.type)
                        .function(ToolCall.Function.builder()
                                .name(toolCall.function.name)
                                .arguments(toolCall.function.arguments)
                                .build())
                        .build());
            }

            log.info("{} call llm stream response {} {}", context.getRequestId(), stringBuilderAll, JSON.toJSONString(toolCalls));

            return ToolCallResponse.builder()
                    .toolCalls(toolCalls)
                    .content(contentAll)
                    .build();
        }

    }


    /**
     * 调用 OpenAI 流式 API（抽象方法，实际实现需要在子类中提供）
     */
    @SneakyThrows
    protected String callOpenAIStream(Map<String, Object> params)
    {
        // 这里是一个简化的流式请求实现示例
        StringBuilder collectedMessages = new StringBuilder();

        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(300, TimeUnit.SECONDS)
                .readTimeout(300, TimeUnit.SECONDS)
                .writeTimeout(300, TimeUnit.SECONDS)
                .build();

        String apiEndpoint = baseUrl + interfaceUrl;

        RequestBody body = RequestBody.create(
                MediaType.parse("application/json"),
                objectMapper.writeValueAsString(params)
        );

        Request.Builder requestBuilder = new Request.Builder()
                .url(apiEndpoint)
                .post(body);

        // 添加适当的认证头
        requestBuilder.addHeader("Authorization", "Bearer " + apiKey);

        Request request = requestBuilder.build();


        try (Response response = client.newCall(request).execute();
             ResponseBody responseBody = response.body()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected response code: " + response);
            }

            if (responseBody != null) {
                String line;

                BufferedReader reader = new BufferedReader(new InputStreamReader(responseBody.byteStream()));
                while ((line = reader.readLine()) != null) {
                    if (line.startsWith("data: ")) {
                        String data = line.substring(6);
                        if (data.equals("[DONE]")) {
                            break;
                        }

                        try {
                            JsonNode chunk = objectMapper.readTree(data);
                            if (chunk.has("choices") && !chunk.get("choices").isEmpty()) {
                                JsonNode choice = chunk.get("choices").get(0);
                                if (choice.has("delta") && choice.get("delta").has("content")) {
                                    String content = choice.get("delta").get("content").asText();
                                    collectedMessages.append(content);
                                    log.info("recv data: {}", content);
                                }
                            }
                        } catch (Exception e) {
                            // 忽略非 JSON 数据
                        }
                    }
                }

                String fullResponse = collectedMessages.toString().trim();
                if (fullResponse.isEmpty()) {
                    throw new IllegalArgumentException("Empty response from streaming LLM");
                } else {
                    return fullResponse;
                }
            } else {
                throw new IOException("Empty response body");
            }
        }

    }

    /**
     * LLM 响应类
     */
    @Data
    @Builder
    @AllArgsConstructor
    @NoArgsConstructor
    public static class ToolCallResponse {
        private String content;
        private List<ToolCall> toolCalls;
        private String finishReason;
        private Integer totalTokens;
        private long duration;
    }

    @Data
    @JsonIgnoreProperties(ignoreUnknown = true)
    public static class OpenAIChoice {
        private Integer index;
        private OpenAIDelta delta;
        private Object logprobs;
        private String finish_reason;
    }

    @Data
    @JsonIgnoreProperties(ignoreUnknown = true)
    public static class OpenAIDelta {
        private String content;
        private List<OpenAIToolCall> tool_calls;
    }

    @Data
    @JsonIgnoreProperties(ignoreUnknown = true)
    public static class OpenAIToolCall {
        private Integer index;
        private String id;
        private String type;
        private OpenAIFunction function;
    }

    @Data
    @JsonIgnoreProperties(ignoreUnknown = true)
    public static class OpenAIFunction {
        private String name;
        private String arguments;
    }

}